from dojo.models import Finding


class DeepfenceThreatmapperMalware:
    def get_findings(self, row, headers, test):
        if "Rule Name" in headers and "Class" in headers:
            return self._parse_old_format(row, headers, test)
        if "Rule Name" in headers and "Node Type" in headers:
            return self._parse_new_format(row, headers, test)
        return None

    def _parse_old_format(self, row, headers, test):
        Rule_Name = row[headers["Rule Name"]]
        Class = row[headers["Class"]]
        File_Name = row[headers["File Name"]]
        Summary = row[headers["Summary"]]
        Severity = row[headers["Severity"]]
        Node_Name = row[headers["Node Name"]]
        NodeType = row[headers["NodeType"]]
        Container_Name = row[headers["Container Name"]]
        Kubernetes_Cluster_Name = row[headers["Kubernetes Cluster Name"]]

        description = (
            f"**Summary:** {Summary}\n"
            f"**Rule Name:** {Rule_Name}\n"
            f"**Class:** {Class}\n"
            f"**File Name:** {File_Name}\n"
            f"**Node Name:** {Node_Name}\n"
            f"**NodeType:** {NodeType}\n"
            f"**Container Name:** {Container_Name}\n"
            f"**Kubernetes Cluster Name:** {Kubernetes_Cluster_Name}\n"
        )

        return Finding(
            title=Rule_Name,
            description=description,
            file_path=File_Name,
            severity=self.severity(Severity),
            static_finding=False,
            dynamic_finding=True,
            test=test,
        )

    def _parse_new_format(self, row, headers, test):
        Rule_Name = row[headers["Rule Name"]]
        File_Name = row[headers["File Name"]]
        Summary = row[headers["Summary"]]
        Severity = row[headers["Severity"]]
        Node_Name = row[headers["Node Name"]]
        Node_Type = row[headers["Node Type"]]
        Kubernetes_Cluster_Name = row[headers["Kubernetes Cluster Name"]]
        Masked = row[headers["Masked"]]

        description = (
            f"**Summary:** {Summary}\n"
            f"**Rule Name:** {Rule_Name}\n"
            f"**File Name:** {File_Name}\n"
            f"**Node Name:** {Node_Name}\n"
            f"**Node Type:** {Node_Type}\n"
            f"**Kubernetes Cluster Name:** {Kubernetes_Cluster_Name}\n"
            f"**Masked:** {Masked}\n"
        )

        return Finding(
            title=Rule_Name,
            description=description,
            file_path=File_Name,
            severity=self.severity(Severity),
            static_finding=False,
            dynamic_finding=True,
            test=test,
        )

    def severity(self, severity_input):
        if severity_input is None:
            return "Info"
        return severity_input.capitalize()
